# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: daniel

import datetime
import os
import matplotlib.pyplot as plt
import pandas as pd
import pathlib
import jinja2
import pytz
import json
import collections
from absl import app, flags

import coin.proto.coin_order_enums_pb2 as coin_enum
from coin.support.proto_log.app.strat_log_replayer import OnLog as strat_log_parser
from coin.support.proto_log.logic.util import (read_strat_log, StratInfo)


FLAGS = flags.FLAGS


def _sumzip(*items):
  return [sum(values) for values in zip(*items)]


def _get_min_max_dt(trading_date):
  return (
      datetime.datetime.combine(trading_date, datetime.time.min.replace(tzinfo=pytz.UTC)),
      datetime.datetime.combine(trading_date, datetime.time.max.replace(tzinfo=pytz.UTC))
  )


class OrderLatencyAnalyzer(object):
  def __init__(self):
    pass

  def retrieve_order_latency_record(self, records):
    dic = dict()
    for record in records:
      if record['og_log']['type'] != "ORDER_EVENT":
        continue
      if record['og_log']['event']['type'] not in ['ORDER_SUBMITTED', 'ORDER_ACCEPTED', 'ORDER_FILLED']:
        continue
      proc_order_id = int(record['og_log']['event']['proc_order_id'])
      if proc_order_id not in dic:
        info = {
            "proc_order_id": proc_order_id,
            "order_submitted_ts": 0,
            "order_accepted_ts": 0,
            "order_filled_ts": 0,
        }
        dic[proc_order_id] = info
      rec = dic[proc_order_id]
      if record['og_log']['event']['type'] == 'ORDER_SUBMITTED':
        rec["order_submitted_ts"] = int(record['og_log']['event']['event_time'])
      elif record['og_log']['event']['type'] == 'ORDER_ACCEPTED':
        rec["order_accepted_ts"] = int(record['og_log']['event']['event_time'])
      elif record['og_log']['event']['type'] == 'ORDER_FILLED':
        rec["order_filled_ts"] = int(record['og_log']['event']['event_time'])

    tmp = []
    for (key, val) in dic.items():
      if val["order_submitted_ts"] == 0:
        continue
      if val["order_accepted_ts"] == 0 and val["order_filled_ts"] == 0:
        continue
      order_accepted_latency = 0
      if val["order_accepted_ts"] != 0:
        order_accepted_latency = val["order_accepted_ts"] - val["order_submitted_ts"]
      order_filled_latency = 0
      if val["order_filled_ts"] != 0:
        order_filled_latency = val["order_filled_ts"] - val["order_submitted_ts"]
      info = {
          "proc_order_id": key,
          "submitted_ts": val["order_submitted_ts"],
          "accepted_latency": order_accepted_latency,
          "filled_latency": order_filled_latency,
      }
      tmp.append(info)

    return sorted(tmp, key=lambda x: x['submitted_ts'])


def dump(root_dir, machine, strategy_name, start_dt, end_dt):
  parser = strat_log_parser(symbol_filter=None, log_type_filter=['OG_LOG'])
  callback = parser.on_log
  strat_info = StratInfo(strategy_name=strategy_name,
                         machine=machine,
                         trading_date=None)
  read_strat_log(strat_info=strat_info,
                 start_time=start_dt,
                 end_time=end_dt,
                 callback=callback,
                 kafka_config=None,
                 root_dir=root_dir)
  return parser.record_list


def process_by_strategy(machine, strategy_name, trading_date, root_dir, output_dir):
  trading_date_tmp = datetime.datetime.strptime(FLAGS.trading_date, '%Y%m%d').date()
  start_dt, end_dt = _get_min_max_dt(trading_date_tmp)

  start_dt_str = start_dt.strftime('%Y%m%dT%H%M%S')
  end_dt_str = end_dt.strftime('%Y%m%dT%H%M%S')
  print('running for %s-%s %s %s ...' % (start_dt_str, end_dt_str, machine, strategy_name))

  records = dump(root_dir, machine, strategy_name, start_dt, end_dt)
  analyzer = OrderLatencyAnalyzer()
  tmp = analyzer.retrieve_order_latency_record(records)
  # save as json file
  output_path = f"{output_dir}/{trading_date}/order_latency"
  if not os.path.exists(output_path):
    os.makedirs(output_path)
  report = f"{output_path}/{strategy_name}_{machine}.json"
  with open(report, 'w') as ofile:
    json.dump(tmp, ofile, indent=2)
  print("writing to %s" % report)


def main(_):
  trading_date = FLAGS.trading_date
  if trading_date is None:
    print("trading date is None")
    exit(1)
  machine_list = None
  if FLAGS.machine is not None:
    machine_list = [elem.strip() for elem in FLAGS.machine.split(',')]
  root_dir = FLAGS.root_dir
  output_dir = FLAGS.output_dir
  for machine in os.listdir(root_dir):
    if machine not in machine_list:
      continue
    fin_path = f"{root_dir}/{machine}/{trading_date}"
    if not os.path.exists(fin_path):
      print('ERROR: path non-exist %s' % fin_path)
      continue
    strategies = []
    for filename in os.listdir(fin_path):
      strategy_name = filename.split('.')[0]
      if strategy_name in strategies:
        continue
      strategies.append(strategy_name)
    for strategy_name in strategies:
      process_by_strategy(machine, strategy_name, trading_date, root_dir, output_dir)


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'root dir specify')
  flags.DEFINE_string('output_dir',
                      '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/og_log_plot',
                      'output_dir.')
  flags.DEFINE_string('tmp_dir',
                      '/tmp',
                      'tmp dir for slurm run')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('local_run', False, "")
  flags.DEFINE_bool('slurm_run', False, "")
  flags.DEFINE_bool('generate_summary_report', False, "generate summary report only")
  app.run(main)
